import codecs

prompt_path = '/home/rickwwang/project_research/fairseq/data-bin/prompt.txt'
all_prompts = codecs.open(prompt_path, 'r', encoding='utf8').readlines()
all_prompts = map(lambda x: x.strip(), all_prompts)

dict_path = '/home/rickwwang/project_research/fairseq/data-bin/writingPrompts-prompt2story_bpe500_fix/dict.prompt.txt'
words = codecs.open(dict_path, 'r', encoding='utf8').readlines()
words = map(lambda x: x.strip().split(' ')[0], words)
prompt_dict = set(words)

# conv-s2s-d0.3_bpe500/40src.txt
# conv-s2s-d0.3-fusion_bpe500/18src.txt
exp_file = []
# exp_file.append('fire_data/running/conv-s2s-d0.3_bpe500/40src.txt')
# exp_file.append('fire_data/running/conv-s2s-d0.3_bpe500/40hyp.txt.1.topk40.t1.0')
exp_file.append('fire_data/running/conv-s2s-d0.3-fusion_bpe500/18src.txt')
exp_file.append('fire_data/running/conv-s2s-d0.3-fusion_bpe500/18hyp.txt.1.topk20.t1.0')


base_path = '/home/rickwwang/project_research/'

story_prompt_path = base_path + exp_file[0]
story_path = base_path + exp_file[1]
result_path = story_path + '.human'
all_stories_prompts = codecs.open(story_prompt_path, 'r', encoding='utf8').readlines()
all_stories = codecs.open(story_path, 'r', encoding='utf8').readlines()
all_stories = zip(all_stories_prompts, all_stories)

promp_to_story = dict()
for s in all_stories:
    key = s[0].strip()
    value = s[1].strip()
    if key in promp_to_story:
        promp_to_story[key].append(value)
    else:
        promp_to_story[key] = [value]

format_stories = []
for p in all_prompts:
    p = p.split(' ')
    p = [w if w in prompt_dict else '<unk>' for w in p]
    p = ' '.join(p)
    ss = promp_to_story[p]

    format_stories.append(p)
    format_stories.append(ss[0])

with codecs.open(result_path, 'w', encoding='utf8') as fout:
    fout.write('\n'.join(format_stories))
    fout.write('\n')